/*
 * Copyright (c) 2020, Peter Abeles. All Rights Reserved.
 *
 * This file is part of BoofCV (http://boofcv.org).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package boofcv.ohos;

import boofcv.alg.color.ColorFormat;
import boofcv.struct.image.GrayF32;
import boofcv.struct.image.GrayU8;
import boofcv.struct.image.ImageBase;
import boofcv.struct.image.ImageDataType;
import boofcv.struct.image.ImageGray;
import boofcv.struct.image.ImageInterleaved;
import boofcv.struct.image.InterleavedF32;
import boofcv.struct.image.InterleavedU8;
import boofcv.struct.image.Planar;
import ohos.media.image.PixelMap;
import ohos.media.image.common.PixelFormat;
import ohos.media.image.common.Size;
import org.ddogleg.struct.DogArray_I8;

import java.nio.ByteBuffer;

import static ohos.media.image.common.PixelFormat.ARGB_8888;

/**
 * Functions for converting ohos Bitmap images into BoofCV formats.  In earlier versions of ohos
 * there is no way to directly access the internal array used by Bitmap.  You have to provide an array for it to
 * be copied into.  This is why the storage array is provided.
 *
 * @author:Peter Abeles
 * @since 2021-05-11
 */
public class ConvertBitmap {

    /**
     * Checks to see if the bitmap is the same shape as the input image. if not a new instance is returned which is
     * otherwise the same instance is returned.
     *
     * @param input
     * @param bitmap
     * @return PixelMap
     */
    public static PixelMap checkDeclare(ImageBase input, PixelMap bitmap) {
        if (bitmap == null) {
            return PixelMap.create(createOpts(input.width, input.height, ARGB_8888));
        } else if (input.width != bitmap.getImageInfo().size.width || input.height != bitmap.getImageInfo().size.height) {
            return PixelMap.create(createOpts(input.width, input.height, bitmap.getImageInfo().pixelFormat));
        } else {
            return bitmap;
        }
    }

    /**
     * resizeStorage
     *
     * @param input
     * @param storage
     * @return DogArray_I8
     */
    public static DogArray_I8 resizeStorage(PixelMap input, DogArray_I8 storage) {
        int byteCount = input.getImageInfo().pixelFormat == ARGB_8888 ? 4 : 2;
        int length = input.getImageInfo().size.width * input.getImageInfo().size.height * byteCount;

        if (storage == null)
            return new DogArray_I8(length);
        else {
            storage.resize(length);
            return storage;
        }
    }

    /**
     * Converts a {@link PixelMap} into a BoofCV image.  Type is determined at runtime.
     *
     * @param input Bitmap image.
     * @param output Output image.  Automatically resized to match input shape.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @param <T>
     * @throws IllegalArgumentException
     */
    public static <T extends ImageBase<T>> void bitmapToBoof(PixelMap input, T output, DogArray_I8 storage) {
        storage = resizeStorage(input, storage);
        if (BOverrideConvertOhos.invokeBitmapToBoof(input, output, storage.data))
            return;

        switch (output.getImageType().getFamily()) {
            case GRAY: {
                if (output.getClass() == GrayF32.class)
                    bitmapToGray(input, (GrayF32) output, storage);
                else if (output.getClass() == GrayU8.class)
                    bitmapToGray(input, (GrayU8) output, storage);
                else
                    throw new IllegalArgumentException("Unsupported BoofCV Image Type");
            }
            break;

            case PLANAR:
                Planar pl = (Planar) output;
                bitmapToPlanar(input, pl, pl.getBandType(), storage);
                break;

            case INTERLEAVED:
                if (output.getClass() == InterleavedU8.class)
                    bitmapToInterleaved(input, (InterleavedU8) output, storage);
                else if (output.getClass() == InterleavedF32.class)
                    bitmapToInterleaved(input, (InterleavedF32) output, storage);
                else
                    throw new IllegalArgumentException("Unsupported BoofCV Image Type");
                break;

            default:
                throw new IllegalArgumentException("Unsupported BoofCV Image Type");
        }
    }

    /**
     * Converts Bitmap image into a single band image of arbitrary type.
     *
     * @param input Input Bitmap image.
     * @param output Output single band image.  If null a new one will be declared.
     * @param imageType Type of single band image.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @param <T>
     * @return The converted gray scale image.
     * @throws IllegalArgumentException
     */
    public static <T extends ImageGray<T>>
    T bitmapToGray(PixelMap input, T output, Class<T> imageType, DogArray_I8 storage) {
        storage = resizeStorage(input, storage);
        if (imageType == GrayF32.class)
            return (T) bitmapToGray(input, (GrayF32) output, storage);
        else if (imageType == GrayU8.class)
            return (T) bitmapToGray(input, (GrayU8) output, storage);
        else
            throw new IllegalArgumentException("Unsupported BoofCV Image Type");
    }

    /**
     * Converts Bitmap image into GrayU8.
     *
     * @param input Input Bitmap image.
     * @param output Output image.  If null a new one will be declared.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @return The converted gray scale image.
     */
    public static GrayU8 bitmapToGray(PixelMap input, GrayU8 output, DogArray_I8 storage) {
        storage = resizeStorage(input, storage);
        if (output == null) {
            output = new GrayU8(input.getImageInfo().size.width, input.getImageInfo().size.height);
        } else {
            output.reshape(input.getImageInfo().size.width, input.getImageInfo().size.height);
        }

        input.readPixels(ByteBuffer.wrap(storage.data));

        ImplConvertBitmap.arrayToGray(storage.data, input.getImageInfo().pixelFormat, output);

        return output;
    }

    /**
     * Converts Bitmap image into GrayF32.
     *
     * @param input Input Bitmap image.
     * @param output Output image.  If null a new one will be declared.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @return The converted gray scale image.
     */
    public static GrayF32 bitmapToGray(PixelMap input, GrayF32 output, DogArray_I8 storage) {
        if (output == null) {
            output = new GrayF32(input.getImageInfo().size.width, input.getImageInfo().size.height);
        } else {
            output.reshape(input.getImageInfo().size.width, input.getImageInfo().size.height);
        }

        storage = resizeStorage(input, storage);
        input.readPixels(ByteBuffer.wrap(storage.data));

        ImplConvertBitmap.arrayToGray(storage.data, input.getImageInfo().pixelFormat, output);

        return output;
    }

    /**
     * Converts Bitmap image into Planar image of the appropriate type.
     *
     * @param input Input Bitmap image.
     * @param output Output image.  If null a new one will be declared.
     * @param type The type of internal single band image used in the Planar image.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @param <T>
     * @return The converted Planar image.
     * @throws IllegalArgumentException
     */
    public static <T extends ImageGray<T>>
    Planar<T> bitmapToPlanar(PixelMap input, Planar<T> output, Class<T> type, DogArray_I8 storage) {
        if (output == null) {
            output = new Planar<>(type, input.getImageInfo().size.width, input.getImageInfo().size.height, 3);
        } else {
            int numBands = Math.min(4, Math.max(3, output.getNumBands()));
            output.reshape(input.getImageInfo().size.width, input.getImageInfo().size.height, numBands);
        }

        storage = resizeStorage(input, storage);
        input.readPixels(ByteBuffer.wrap(storage.data));

        if (type == GrayU8.class)
            ImplConvertBitmap.arrayToPlanar_U8(storage.data, input.getImageInfo().pixelFormat, (Planar) output);
        else if (type == GrayF32.class)
            ImplConvertBitmap.arrayToPlanar_F32(storage.data, input.getImageInfo().pixelFormat, (Planar) output);
        else
            throw new IllegalArgumentException("Unsupported BoofCV Type");

        return output;
    }

    /**
     * Converts Bitmap image into InterleavedU8.
     *
     * @param input Input Bitmap image.
     * @param output Output image.  If null a new one will be declared.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @return The converted gray scale image.
     */
    public static InterleavedU8 bitmapToInterleaved(PixelMap input, InterleavedU8 output, DogArray_I8 storage) {
        if (output == null) {
            output = new InterleavedU8(input.getImageInfo().size.width, input.getImageInfo().size.height, 3);
        } else {
            if (output.getNumBands() < 3 || output.getNumBands() > 4)
                output.reshape(input.getImageInfo().size.width, input.getImageInfo().size.height, 3);
            else
                output.reshape(input.getImageInfo().size.width, input.getImageInfo().size.height);
        }

        storage = resizeStorage(input, storage);
        input.readPixels(ByteBuffer.wrap(storage.data));

        ImplConvertBitmap.arrayToInterleaved_U8(storage.data, input.getImageInfo().pixelFormat, output);

        return output;
    }

    /**
     * Converts Bitmap image into 	public static InterleavedF32 bitmapToInterleaved( Bitmap input, InterleavedU8 output, @Nullable DogArray_I8 storage ) {.
     *
     * @param input Input Bitmap image.
     * @param output Output image.  If null a new one will be declared.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @return The converted gray scale image.
     */
    public static InterleavedF32 bitmapToInterleaved(PixelMap input, InterleavedF32 output, DogArray_I8 storage) {
        if (output == null) {
            output = new InterleavedF32(input.getImageInfo().size.width, input.getImageInfo().size.height, 3);
        } else {
            if (output.getNumBands() < 3 || output.getNumBands() > 4)
                output.reshape(input.getImageInfo().size.width, input.getImageInfo().size.height, 3);
            else
                output.reshape(input.getImageInfo().size.width, input.getImageInfo().size.height);
        }

        storage = resizeStorage(input, storage);
        input.readPixels(ByteBuffer.wrap(storage.data));

        ImplConvertBitmap.arrayToInterleaved_F32(storage.data, input.getImageInfo().pixelFormat, output);

        return output;
    }

    /**
     * Converts many BoofCV image types into a Bitmap.
     *
     * @param input Input BoofCV image.
     * @param output Output Bitmap image.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @throws IllegalArgumentException
     */
    public static void boofToBitmap(ImageBase input, PixelMap output, DogArray_I8 storage) {
        if (BOverrideConvertOhos.invokeBoofToBitmap(ColorFormat.RGB, input, output, storage.data))
            return;

        if (input instanceof Planar) {
            planarToBitmap((Planar) input, output, storage);
        } else if (input instanceof ImageGray) {
            grayToBitmap((ImageGray) input, output, storage);
        } else if (input instanceof ImageInterleaved) {
            interleavedToBitmap((ImageInterleaved) input, output, storage);
        } else {
            throw new IllegalArgumentException("Unsupported input image type");
        }
    }

    /**
     * boofToBitmap
     *
     * @param color
     * @param input
     * @param output
     * @param storage
     * @throws IllegalArgumentException
     */
    public static void boofToBitmap(ColorFormat color, ImageBase input, PixelMap output, DogArray_I8 storage) {
        storage = resizeStorage(output, storage);
        if (BOverrideConvertOhos.invokeBoofToBitmap(color, input, output, storage.data))
            return;

        if (input instanceof ImageGray) {
            grayToBitmap((ImageGray) input, output, storage);
            return;
        }

        switch (color) {
            case RGB: {
                boofToBitmap(input, output, storage);
            }
            return;

            case YUV: {
                if (input instanceof ImageInterleaved) {
                    interleavedYuvToBitmap((ImageInterleaved) input, output, storage);
                    return;
                }
            }
            break;

            default:
                break;
        }
        throw new IllegalArgumentException("Unsupported input image type");
    }

    /**
     * Converts ImageGray into Bitmap.
     *
     * @param input Input gray scale image.
     * @param output OutputPixelMapimage.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @throws IllegalArgumentException
     */
    public static void grayToBitmap(ImageGray input, PixelMap output, DogArray_I8 storage) {
        if (input instanceof GrayU8)
            grayToBitmap((GrayU8) input, output, storage);
        else if (input instanceof GrayF32)
            grayToBitmap((GrayF32) input, output, storage);
        else
            throw new IllegalArgumentException("Unsupported BoofCV Type: " + input);
    }

    /**
     * Converts ImageGray into Bitmap.
     *
     * @param input Input gray scale image.
     * @param output Output PixelMap image.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @throws IllegalArgumentException
     */
    public static void grayToBitmap(GrayU8 input, PixelMap output, DogArray_I8 storage) {
        if (output.getImageInfo().size.width != input.getWidth() || output.getImageInfo().size.height != input.getHeight()) {
            throw new IllegalArgumentException("Image shapes are not the same");
        }

        storage = resizeStorage(output, storage);

        ImplConvertBitmap.grayToArray(input, storage.data, output.getImageInfo().pixelFormat);
        output.writePixels(ByteBuffer.wrap(storage.data));
    }

    /**
     * Converts ImageGray into Bitmap.
     *
     * @param input Input gray scale image.
     * @param output Output PixelMap image.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @throws IllegalArgumentException
     */
    public static void grayToBitmap(GrayF32 input, PixelMap output, DogArray_I8 storage) {
        if (output.getImageInfo().size.width != input.getWidth() || output.getImageInfo().size.height != input.getHeight()) {
            throw new IllegalArgumentException("Image shapes are not the same");
        }

        storage = resizeStorage(output, storage);

        ImplConvertBitmap.grayToArray(input, storage.data, output.getImageInfo().pixelFormat);
        output.writePixels(ByteBuffer.wrap(storage.data));
    }

    /**
     * Converts Planar image into Bitmap.
     *
     * @param input Input Planar image.
     * @param output Output PixelMap image.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @param <T>
     * @throws IllegalArgumentException
     */
    public static <T extends ImageGray<T>>
    void planarToBitmap(Planar<T> input, PixelMap output, DogArray_I8 storage) {
        if (output.getImageInfo().size.width != input.getWidth() || output.getImageInfo().size.height != input.getHeight()) {
            throw new IllegalArgumentException("Image shapes are not the same");
        }

        storage = resizeStorage(output, storage);

        if (input.getBandType() == GrayU8.class)
            ImplConvertBitmap.planarToArray_U8((Planar) input, storage.data, output.getImageInfo().pixelFormat);
        else if (input.getBandType() == GrayF32.class)
            ImplConvertBitmap.planarToArray_F32((Planar) input, storage.data, output.getImageInfo().pixelFormat);
        else
            throw new IllegalArgumentException("Unsupported BoofCV Type");
        output.writePixels(ByteBuffer.wrap(storage.data));
    }

    /**
     * Converts {@link ImageInterleaved} image into Bitmap.
     *
     * @param input Input Planar image.
     * @param output Output PixelMap image.
     * @param storage Byte array used for internal storage. If null it will be declared internally.
     * @param <T>
     * @throws IllegalArgumentException
     */
    public static <T extends ImageInterleaved<T>>
    void interleavedToBitmap(T input, PixelMap output, DogArray_I8 storage) {
        if (output.getImageInfo().size.width != input.getWidth() || output.getImageInfo().size.height != input.getHeight()) {
            throw new IllegalArgumentException("Image shapes are not the same");
        }

        storage = resizeStorage(output, storage);

        if (input.getImageType().getDataType() == ImageDataType.U8)
            ImplConvertBitmap.interleavedToArray((InterleavedU8) input, storage.data, output.getImageInfo().pixelFormat);
        else if (input.getImageType().getDataType() == ImageDataType.F32)
            ImplConvertBitmap.interleavedToArray((InterleavedF32) input, storage.data, output.getImageInfo().pixelFormat);
        else
            throw new IllegalArgumentException("Unsupported BoofCV Type");
        output.writePixels(ByteBuffer.wrap(storage.data));
    }

    /**
     * interleavedYuvToBitmap
     *
     * @param input
     * @param output
     * @param storage
     * @param <T>
     * @throws IllegalArgumentException
     */
    public static <T extends ImageInterleaved<T>>
    void interleavedYuvToBitmap(T input, PixelMap output, DogArray_I8 storage) {
        if (output.getImageInfo().size.width != input.getWidth() || output.getImageInfo().size.height != input.getHeight()) {
            throw new IllegalArgumentException("Image shapes are not the same");
        }

        storage = resizeStorage(output, storage);

        if (input.getImageType().getDataType() == ImageDataType.U8) {
            switch (output.getImageInfo().pixelFormat) {
                case ARGB_8888:
                    ImplConvertBitmap.interleavedYuvToArgb8888((InterleavedU8) input, storage.data);
                    output.writePixels(ByteBuffer.wrap(storage.data));
                    return;
                case RGB_565:
                    ImplConvertBitmap.interleavedYuvToRGB565((InterleavedU8) input, storage.data);
                    output.writePixels(ByteBuffer.wrap(storage.data));
                    return;

                default:
                    break;
            }
        }
        throw new IllegalArgumentException("Unsupported BoofCV Type");
    }

    /**
     * Converts GrayU8 into a new Bitmap.
     *
     * @param input Input gray scale image.
     * @param config Type ofPixelMapimage to create.
     * @return PixelMap
     */
    public static PixelMap grayToBitmap(GrayU8 input, PixelFormat config) {
        PixelMap output = PixelMap.create(createOpts(input.width, input.height, config));

        grayToBitmap(input, output, null);

        return output;
    }

    /**
     * Converts GrayF32 into a new Bitmap.
     *
     * @param input Input gray scale image.
     * @param config Type ofPixelMapimage to create.
     * @return grayToBitmap
     */
    public static PixelMap grayToBitmap(GrayF32 input, PixelFormat config) {
        PixelMap output = PixelMap.create(createOpts(input.width, input.height, config));

        grayToBitmap(input, output, null);

        return output;
    }

    public static PixelMap.InitializationOptions createOpts(int width, int height, PixelFormat pixelFormat) {
        PixelMap.InitializationOptions initOptions = new PixelMap.InitializationOptions();
        initOptions.pixelFormat = pixelFormat;
        initOptions.size = new Size(width, height);
        initOptions.editable = true;
        return initOptions;
    }
}
